import numpy as np
import meshio
import os

class MeshioWriter:
    def __init__(self, reader):
        """
        Initialize the MeshioWriter.

        :param reader: BinReader object.

        """
        self.reader = reader

    def create_mesh(self):
        """
        Create a meshio Mesh with points and cells.

        :return: meshio.Mesh

        """
        # Process points
        GRID_VTX_array = {}
        if self.reader.get_static_data("GRID_VTX", GRID_VTX_array):
            data = np.array(GRID_VTX_array["data"], dtype=np.float64)
            points = data.reshape(-1, 3)
            print(f"Found {len(points)} points.")
        else:
            raise ValueError("GRID_VTX data not found.")

        # Process cells
        NODE_LIST_array = {}
        if self.reader.get_static_data("NODE_LIST", NODE_LIST_array):
            cell_data = np.array(NODE_LIST_array["data"], dtype=np.int64) - 1
            n_cells = len(cell_data) // 8
            cell_data = cell_data.reshape(-1, 8)
            cells = [("hexahedron", cell_data)]

        else:
            raise ValueError("NODE_LIST data not found.")

        return meshio.Mesh(points, cells)

    def write_xdmf(self, output_folder, filename):
        os.makedirs(output_folder, exist_ok=True)
        mesh = self.create_mesh()

        # Write the main mesh file (contains topology and geometry)
        main_mesh_filename = f"{filename}_mesh.xdmf"
        main_mesh_filepath = os.path.join(output_folder, main_mesh_filename)
        meshio.write(main_mesh_filepath, mesh, file_format="xdmf")
        print(f"Successfully wrote main mesh to {main_mesh_filepath}")

        # Prepare cell data
        cell_data = {}
        for info in self.reader.stat_props:
            prop_array = {}
            if self.reader.get_static_data(info.name, prop_array):
                if prop_array["on_grid"]:
                    if prop_array["type"] in ["d", "i"]:
                        print(
                            f"Adding property {info.name} for static to mesh"
                        )
                        cell_data[info.name] = [prop_array["data"]]
                    else:
                        print(
                            f"Unknown data type {prop_array['type']} for property {info.name}"
                        )
                else:
                    print(f"Property {info.name} is not on grid.")
            else:
                print(f"Property {info.name} not found for static.")

        if cell_data:
            static_filename = f"{filename}_static.xdmf"
            static_filepath = os.path.join(output_folder, static_filename)

            # Create a new mesh with only the static data
            # mesh.cell_data = cell_data
            # meshio.write(static_filepath, mesh, file_format="xdmf")
            static_mesh = meshio.Mesh(points=mesh.points, cells=mesh.cells, cell_data=cell_data)
            static_mesh.write(static_filepath, file_format="xdmf")
            print(f"Successfully wrote static data to {static_filepath}")

        for ind, dynamic_infos in enumerate(self.reader.dym_props):
            tstep = int(dynamic_infos[0].time)
            if ind >= len(self.reader.dym_props):
                print(f"Time step {tstep} does not exist. Skipping.")
                continue

            # Prepare cell data
            cell_data = {}
            for info in dynamic_infos:
                prop_array = {}
                if self.reader.get_dynamic_data(ind, info.name, prop_array):
                    if prop_array["on_grid"]:
                        if prop_array["type"] in ["d", "i"]:
                            print(
                                f"Adding property {info.name} for timestep {tstep} to mesh"
                            )
                            cell_data[info.name] = [prop_array["data"]]
                        else:
                            print(
                                f"Unknown data type {prop_array['type']} for property {info.name}"
                            )
                    else:
                        print(f"Property {info.name} is not on grid.")
                else:
                    print(f"Property {info.name} not found for timestep {tstep}.")

            if cell_data:
                timestep_filename = f"{filename}_t{tstep}.xdmf"
                timestep_filepath = os.path.join(output_folder, timestep_filename)

                # mesh.cell_data = cell_data
                # meshio.write(timestep_filepath, mesh, file_format="xdmf")
                dynamic_mesh = meshio.Mesh(points=mesh.points, cells=mesh.cells, cell_data=cell_data)
                dynamic_mesh.write(timestep_filepath, file_format="xdmf")
                print(f"Successfully wrote mesh for timestep {tstep} to {timestep_filepath}")

class WellioWriter:
    def __init__(self, reader):
        """
        Initialize the MeshioWriter.

        :param reader: WELLReader object.

        """
        self.reader = reader
        self.AxisOffset = np.zeros(3, dtype=np.float64)
        self.well_index_map = {}

    def well_offset(self, offset_well, offset_mesh):
        """
        Refresh the offset of well data.

        :param offset_well: Offset of well data.
        :param offset_mesh: Offset of mesh data.

        """
        self.AxisOffset = offset_well - offset_mesh
    
    def create_well_mesh(self):
        """
        Create a meshio Mesh with points and cells.

        :return: meshio.Mesh
        """
        # Process points
        points = []
        lines = []
        coord_names = ["XCOORD", "YCOORD", "DEPTH"]

        for well_num in self.reader.WELL_num:
            coord_array = {}
            for name in coord_names:
                if not self.reader.get_static_data(well_num, name, coord_array) or name not in coord_array:
                    raise ValueError(f"{name} data not found for WELL {well_num}.")

            x_data, y_data, z_data = (np.array(coord_array[name], dtype=np.float64) for name in coord_names)

            if not (len(x_data) == len(y_data) == len(z_data)):
                raise ValueError(f"Coordinate data length mismatch for WELL {well_num}.")
            if not np.all(np.isfinite([x_data, y_data, z_data])):
                raise ValueError(f"Invalid coordinate data for WELL {well_num}.")

            start_index = len(points)  # Starting index for this well's points
            
            offset = self.AxisOffset 
            points.extend([[x + offset[0], y + offset[1], z + offset[2]] for x, y, z in zip(x_data, y_data, z_data)])

            lines.extend([[start_index + i, start_index + i + 1] for i in range(len(x_data) - 1)])

            end_index = len(points) - 1  # Ending index for this well's points
            self.well_index_map[well_num] = (start_index, end_index)

        points = np.array(points, dtype=np.float64)
        lines = np.array(lines, dtype=np.int64)
        cells = [("line", lines)]

        print(f"Generated {len(points)} points and {len(lines)} lines.")

        return meshio.Mesh(points, cells)
    
    def write_xdmf(self, output_folder, filename):
        """
        Write XDMF files for all timesteps.

        :param output_folder: Directory path for output files.
        :param filename: Base filename for output files.
        """
        os.makedirs(output_folder, exist_ok=True)
        mesh = self.create_well_mesh()

        for ind, dynamic_infos in enumerate(self.reader.WELL_dym_props.values()):             
            tstep = self.reader.WELL_timestep[ind]
            if tstep not in self.reader.WELL_dym_props:
                print(f"Time step {tstep} does not exist. Skipping.")
                continue

            point_data = {}
            for well_num, props in dynamic_infos.items():
                if well_num not in self.well_index_map:
                    print(f"WELL {well_num} does not exist in the global mesh. Skipping.")
                    continue

                start_index, end_index = self.well_index_map[well_num]
                num_points = end_index - start_index + 1

                for prop_name, *_ in props:
                    prop_array = {}
                    if self.reader.get_dynamic_data(tstep, well_num, prop_name, prop_array):
                        prop_values = np.array(prop_array[prop_name], dtype=np.float64)

                        if len(prop_values) != num_points:
                            raise ValueError(
                                f"Property {prop_name} for WELL {well_num} does not match point count "
                                f"({len(prop_values)} != {num_points})."
                            )
                        
                        if prop_name not in point_data:
                            point_data[prop_name] = np.zeros(len(mesh.points), dtype=np.float64)
                        point_data[prop_name][start_index:end_index + 1] = prop_values

                        print(f"WELL {well_num}: Adding property {prop_name} for timestep {tstep}.")
                    else:
                        print(f"WELL {well_num}: Property {prop_name} not found for timestep {tstep}.")

            if point_data:
                mesh.point_data = point_data
                timestep_filename = f"{filename}_well_t{int(tstep)}.xdmf"
                output_filepath = os.path.join(output_folder, timestep_filename)
                try:
                    meshio.write(output_filepath, mesh, file_format="xdmf")
                    print(f"Successfully wrote mesh for timestep {tstep} to {output_filepath}")
                except Exception as e:
                    print(f"Failed to write mesh for timestep {tstep}. Error: {e}")
            else:
                print(f"No data to write for timestep {tstep}.")